#include <vector>

#include <gtest/gtest.h>
#include <gmock/gmock.h>

#define protected public
#define private public
#include "single_op/op_executor.h"
#include "executor/stream_executor.h"
#undef private
#undef protected

#define protected public
#define private public
#include "single_op/op_model_manager.h"
#include "single_op/shape_range_utils.h"
#undef private
#undef protected

#define protected public
#define private public
#include "utils/acl_op_map.h"
#undef private
#undef protected

#include "acl/acl.h"
#include "runtime/rt.h"
#include "utils/file_utils.h"
#include "single_op/op_model_parser.h"
#include "single_op/executor/resource_manager.h"
#include "utils/math_utils.h"
#include "common/ge_types.h"
#include "acl_stub.h"

using namespace std;
using namespace testing;
using namespace acl;
using namespace ge;

static aclError SelectAclopBatchNorm(int numInputs,
                            const aclTensorDesc *const inputDesc[],
                            int numOutputs,
                            const aclTensorDesc *const outputDesc[],
                            const aclopAttr *opAttr,
                            aclopKernelDesc *aclopKernelDes)
{
    aclopKernelDes->kernelId = "kernel1";
    return ACL_SUCCESS;
}

class OpExecutorTest : public testing::Test {
protected:
    void SetUp()
    {
        int64_t shape[] = {16, 16};
        inputs_[0] = aclCreateDataBuffer((void *) 0x12345, 1024);
        inputs_[1] = aclCreateDataBuffer((void *) 0x12345, 1024);
        outputs_[0] = aclCreateDataBuffer((void *) 0x12345, 1024);
        input_desc_[0] = aclCreateTensorDesc(ACL_FLOAT, 2, shape, ACL_FORMAT_ND);
        input_desc_[1] = aclCreateTensorDesc(ACL_FLOAT, 2, shape, ACL_FORMAT_ND);
        input_desc_[2] = nullptr;
        output_desc_[0] = aclCreateTensorDesc(ACL_FLOAT, 2, shape, ACL_FORMAT_ND);
        output_desc_[1] = nullptr;
    }

    void TearDown()
    {
        Mock::VerifyAndClear((void *)(&MockFunctionTest::aclStubInstance()));
        aclDestroyDataBuffer(inputs_[0]);
        aclDestroyDataBuffer(inputs_[1]);
        aclDestroyDataBuffer(outputs_[0]);
        aclDestroyTensorDesc(input_desc_[0]);
        aclDestroyTensorDesc(input_desc_[1]);
        aclDestroyTensorDesc(output_desc_[0]);
    }

    const aclDataBuffer *inputs_[2];
    aclDataBuffer *outputs_[1];
    const aclTensorDesc *input_desc_[3];
    const aclTensorDesc *output_desc_[2];
};

Status LoadSingleOpMock(const std::string &modelName,
                        const ge::ModelData &modelData,
                        void *stream,
                        SingleOp **single_op,
                        const uint64_t model_id)
{
    *single_op = (SingleOp *)0x12345678;
    return SUCCESS;
}

Status LoadDynamicSingleOpMock(const std::string &modelName,
                                const ge::ModelData &modelData,
                                void *stream,
                                DynamicSingleOp **single_op,
                                const uint64_t model_id)
{
    *single_op = (DynamicSingleOp *)0x12345678;
    return SUCCESS;
}

TEST_F(OpExecutorTest, TestLoadSingleOp)
{
    EXPECT_CALL(MockFunctionTest::aclStubInstance(), LoadSingleOpV2(_, _,_,_,_))
        .WillOnce(Return(RT_FAILED))
        .WillRepeatedly(Invoke(LoadSingleOpMock));

    EXPECT_CALL(MockFunctionTest::aclStubInstance(), LoadDynamicSingleOpV2(_, _,_,_,_))
        .WillOnce(Return(RT_FAILED))
        .WillRepeatedly(Invoke(LoadDynamicSingleOpMock));

    OpModel opModel;
    auto aclStream = (aclrtStream) 0x1234;
    SingleOp *nullSingleOp = nullptr;
    DynamicSingleOp *nullDynamicSingleOp = nullptr;
    auto *expectedSingleOp = (SingleOp *) 0x12345678;
    auto *expectedDynamicSingleOp = (DynamicSingleOp *) 0x12345678;

    SingleOp *SingleOp = nullptr;
    // first call will return failed
    ASSERT_NE(OpExecutor::LoadSingleOp(opModel, aclStream, &SingleOp), ACL_SUCCESS);
    // other call will return success
    ASSERT_EQ(OpExecutor::LoadSingleOp(opModel, aclStream, &SingleOp), ACL_SUCCESS);
    ASSERT_EQ(SingleOp, expectedSingleOp);

    DynamicSingleOp *dynamicSingleOp = nullptr;
    ASSERT_NE(OpExecutor::LoadDynamicSingleOp(opModel, aclStream, &dynamicSingleOp), ACL_SUCCESS);
    ASSERT_EQ(OpExecutor::LoadDynamicSingleOp(opModel, aclStream, &dynamicSingleOp), ACL_SUCCESS);
    ASSERT_EQ(dynamicSingleOp, expectedDynamicSingleOp);
}

TEST_F(OpExecutorTest, DoExecuteAsyncTest)
{
    auto *singleOp = (SingleOp *) 0x12345678;
    EXPECT_CALL(MockFunctionTest::aclStubInstance(), ExecuteAsync(_,_,_))
        .WillOnce(Return(PARAM_INVALID))
        .WillRepeatedly(Return(SUCCESS));

    AclOp aclOp;
    aclOp.opType = "Add";
    aclOp.numInputs = 2;
    aclOp.numOutputs = 1;
    aclTensorDesc *inputDesc[2];
    aclTensorDesc *outputDesc[1];
    int64_t shape[]{16, 16};
    inputDesc[0] = aclCreateTensorDesc(ACL_FLOAT16, 2, shape, ACL_FORMAT_ND);
    inputDesc[1] = aclCreateTensorDesc(ACL_FLOAT16, 2, shape, ACL_FORMAT_ND);
    outputDesc[0] = aclCreateTensorDesc(ACL_FLOAT16, 2, shape, ACL_FORMAT_ND);
    aclOp.inputDesc = inputDesc;
    aclOp.outputDesc = outputDesc;
    bool executeWithExactModel = true;
    ASSERT_NE(OpExecutor::DoExecuteAsync(singleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);
    ASSERT_EQ(OpExecutor::DoExecuteAsync(singleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);

    aclSetTensorPlaceMent(inputDesc[0], aclMemType::ACL_MEMTYPE_HOST);
    ASSERT_EQ(OpExecutor::DoExecuteAsync(singleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);

    aclSetTensorPlaceMent(outputDesc[0], aclMemType::ACL_MEMTYPE_HOST);
    ASSERT_EQ(OpExecutor::DoExecuteAsync(singleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);

    aclDestroyTensorDesc(inputDesc[0]);
    aclDestroyTensorDesc(inputDesc[1]);
    aclDestroyTensorDesc(outputDesc[0]);
}

TEST_F(OpExecutorTest, CheckConstTensor)
{
    int64_t shape[]{16, -1};
    aclTensorDesc *desc = aclCreateTensorDesc(ACL_FLOAT16, 2, shape, ACL_FORMAT_ND);
    desc->isConst = true;
    ASSERT_EQ(desc->CheckConstTensor(false), true);
    desc->isConst = false;
    ASSERT_EQ(desc->CheckConstTensor(true), false);
    desc->memtype = ACL_MEMTYPE_HOST;
    ASSERT_EQ(desc->CheckConstTensor(true), true);
    aclDestroyTensorDesc(desc);
}

TEST_F(OpExecutorTest, DoExecuteAsyncDynamicSuccessTest)
{
    auto *dynamicSingleOp = (DynamicSingleOp *) 0x12345678;
    EXPECT_CALL(MockFunctionTest::aclStubInstance(), ExecuteAsync(_,_,_,_,_))
        .WillOnce(Return(PARAM_INVALID))
        .WillRepeatedly(Return(SUCCESS));
    AclOp aclOp;
    aclOp.opType = "Add";
    aclOp.numInputs = 2;
    aclOp.numOutputs = 1;
    int64_t shape[]{16, -1};
    ge::SmallVector<int64_t, ge::kDefaultMaxInputNum> storageDims;
    storageDims.push_back(16);
    storageDims.push_back(-1);
    aclopAttr *opAttr = aclopCreateAttr();
    aclTensorDesc *inputDesc[2];
    aclTensorDesc *outputDesc[1];
    inputDesc[0] = aclCreateTensorDesc(ACL_FLOAT16, 2, shape, ACL_FORMAT_ND);
    const_cast<aclTensorDesc *>(inputDesc[0])->storageDims = storageDims;
    inputDesc[0]->IsDynamicTensor();
    inputDesc[1] = aclCreateTensorDesc(ACL_FLOAT16, 2, shape, ACL_FORMAT_ND);
    const_cast<aclTensorDesc *>(inputDesc[1])->storageDims = storageDims;
    outputDesc[0] = aclCreateTensorDesc(ACL_FLOAT16, 2, shape, ACL_FORMAT_ND);
    const_cast<aclTensorDesc *>(outputDesc[0])->storageDims = storageDims;
    aclOp.inputDesc = inputDesc;
    aclOp.outputDesc = outputDesc;
    aclOp.opAttr = opAttr;
    ASSERT_NE(OpExecutor::DoExecuteAsync(dynamicSingleOp, aclOp, inputs_, outputs_), ACL_SUCCESS);
    ASSERT_EQ(OpExecutor::DoExecuteAsync(dynamicSingleOp, aclOp, inputs_, outputs_), ACL_SUCCESS);
    
    bool executeWithExactModel = true;   
    aclSetTensorStorageFormat(const_cast<aclTensorDesc *>(aclOp.inputDesc[0]), ACL_FORMAT_ND);
    ASSERT_EQ(OpExecutor::DoExecuteAsync(dynamicSingleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);

    aclSetTensorStorageFormat(const_cast<aclTensorDesc *>(aclOp.outputDesc[0]), ACL_FORMAT_ND);
    ASSERT_EQ(OpExecutor::DoExecuteAsync(dynamicSingleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);

    aclOp.exeucteType = ACL_OP_EXECUTE_V2;
    ASSERT_EQ(OpExecutor::DoExecuteAsync(dynamicSingleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);

    aclSetTensorPlaceMent(inputDesc[0], aclMemType::ACL_MEMTYPE_HOST);
    ASSERT_EQ(OpExecutor::DoExecuteAsync(dynamicSingleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);

    aclSetTensorPlaceMent(outputDesc[0], aclMemType::ACL_MEMTYPE_HOST);
    ASSERT_EQ(OpExecutor::DoExecuteAsync(dynamicSingleOp, aclOp, inputs_, outputs_, executeWithExactModel), ACL_SUCCESS);

    aclDestroyTensorDesc(inputDesc[0]);
    aclDestroyTensorDesc(inputDesc[1]);
    aclDestroyTensorDesc(outputDesc[0]);
    aclopDestroyAttr(opAttr);
}

TEST_F(OpExecutorTest, TestCaseExecuteAsync)
{
    aclopUnregisterCompileFunc("BatchNorm");
    aclopRegisterCompileFunc("BatchNorm", SelectAclopBatchNorm);
    ASSERT_EQ(aclopCreateKernel("BatchNorm", "kernel1", "kernel1", (void *)0x1000, 1024, ACL_ENGINE_AICORE, nullptr),
              ACL_SUCCESS);
    ASSERT_EQ(aclopUpdateParams("BatchNorm", 1, input_desc_, 1, output_desc_, nullptr),
	      ACL_SUCCESS);
    int res1 = 0;
    size_t res2 = 0;
    ASSERT_NE(CheckIntAddOverflow(2147483647, 1, res1), ACL_SUCCESS);
    ASSERT_NE(CheckSizeTAddOverflow(0xffffffffffffffff, 1, res2), ACL_SUCCESS);
    ASSERT_NE(CheckSizeTMultiOverflow(0xffffffffffffffff, 2, res2), ACL_SUCCESS);
    aclopUnregisterCompileFunc("BatchNorm");
}

TEST_F(OpExecutorTest, TestResourceManager)
{
    ResourceManager resource_manager((void*)0x1234);
    const char* str = "str";
    ASSERT_EQ(resource_manager.GetMemory((void**)&str, 10),
               ACL_SUCCESS);
}

TEST_F(OpExecutorTest, TestInitTbeTask)
{
    OpKernelDesc desc;
    ResourceManager *mng = new ResourceManager((void *)0x1234);
    TbeOpTask task(nullptr, 0);
    desc.workspaceSizes.resize(18);
    StreamExecutor se(mng, nullptr);
    ASSERT_EQ(se.InitTbeTask(desc, 0, 0, task), ACL_ERROR_INVALID_PARAM);
    desc.workspaceSizes.clear();
    desc.workspaceSizes.push_back(32);
    ASSERT_EQ(se.InitTbeTask(desc, 0, 0, task), ACL_SUCCESS);
}

TEST_F(OpExecutorTest, TestAllocateWorkspaces)
{
    std::vector<size_t> workspaceSizes;
    vector<uintptr_t> workspaces;
    workspaceSizes.resize(18);
    StreamExecutor se(nullptr, nullptr);
    se.AllocateWorkspaces(workspaceSizes, workspaces);
}

TEST_F(OpExecutorTest, TestExecutors)
{
    StreamExecutor * exec = Executors::GetOrCreate(nullptr, nullptr);
    ASSERT_NE(exec, nullptr);
    ASSERT_EQ(exec, Executors::GetOrCreate(nullptr, nullptr));
    Executors::RemoveExecutor(nullptr, nullptr);
}

TEST_F(OpExecutorTest, StreamExecuteAsyncTest)
{
    AclOp aclOpDesc;
    aclDataBuffer *const *inputs;
    aclDataBuffer *const *outputs;
    ResourceManager *mng = new ResourceManager((void *)0x1234);
    StreamExecutor executor(mng, nullptr);
    aclError ret = executor.ExecuteAsync(aclOpDesc, inputs, outputs);
    EXPECT_EQ(ret, ACL_ERROR_OP_NOT_FOUND);
}

TEST_F(OpExecutorTest, StreamExecuteAsyncTest1)
{
    AclOp aclOpDesc;
    aclDataBuffer *const inputs[2]{0};
    aclDataBuffer *const outputs[2]{0};
    aclOpDesc.isMatched = true;
    const aclrtStream stream = (aclrtStream) 0x1234;
    auto ret = OpExecutor::ExecuteAsync(aclOpDesc, inputs, outputs, stream);
    EXPECT_EQ(ret, ACL_SUCCESS);

    aclOpDesc.isDynamic = true;
    ret = OpExecutor::ExecuteAsync(aclOpDesc, inputs, outputs, stream);
    EXPECT_EQ(ret, ACL_SUCCESS);
}

TEST_F(OpExecutorTest, StreamExecuteAsyncTest02)
{
    OpKernelDesc kernelDesc;
    size_t i = 0;
    for (size_t i = 0; i < 18; ++i) {
        kernelDesc.workspaceSizes.push_back(1);
    }

    int32_t numInputs = 1;
    int32_t numOutputs = 1;
    aclDataBuffer *const *inputs;
    aclDataBuffer *const *outputs;
    ResourceManager *mng = new ResourceManager((void *)0x1234);
    StreamExecutor executor(mng, nullptr);
    aclError ret = executor.ExecuteAsync(kernelDesc, numInputs, inputs, numOutputs, outputs);
    EXPECT_EQ(ret, ACL_ERROR_INVALID_PARAM);
}


TEST_F(OpExecutorTest, TestStaticExecuteAsyncWithOpHandle)
{
    EXPECT_CALL(MockFunctionTest::aclStubInstance(), LoadSingleOpV2(_,_,_,_,_))
        .WillOnce(Return(ACL_ERROR_GE_LOAD_MODEL))
        .WillRepeatedly(Return(SUCCESS));
    
    EXPECT_CALL(MockFunctionTest::aclStubInstance(), ExecuteAsync(_,_,_,_,_))
        .WillRepeatedly(Return(SUCCESS));

    // aclOp.isCompile = false;
    aclopHandle *handle = nullptr;

    OpModelDef modelDef;
    modelDef.opType = "AAA";
    modelDef.modelPath = "BBB";

    int64_t shape[] = {16, 16};
    modelDef.inputDescArr.emplace_back(ACL_FLOAT, 2, shape, ACL_FORMAT_ND);
    modelDef.inputDescArr.emplace_back(ACL_FLOAT, 2, shape, ACL_FORMAT_ND);
    modelDef.outputDescArr.emplace_back(ACL_FLOAT, 2, shape, ACL_FORMAT_ND);

    OpModel opModel;

    char_t *data = new char_t(32);
    const std::shared_ptr<void> p_data(data, [](const char_t *const p) { delete[]p; });
    opModel.data = p_data;
    opModel.size = 32U;
    opModel.name = "path";

    auto &instance = OpModelManager::GetInstance();
    instance.modelCache_.Add(modelDef, opModel);
    EXPECT_EQ(instance.RegisterModel(std::move(modelDef), OpModelManager::GetInstance().opModels_, false), ACL_SUCCESS);

    aclopAttr *attr = aclopCreateAttr();
    aclopCreateHandle("AAA", 2, input_desc_, 1, output_desc_, attr, &handle);
    const aclrtStream stream = (aclrtStream) 0x1234;
    EXPECT_NE(aclopExecWithHandle(handle, 2, inputs_, 1, outputs_, stream), ACL_SUCCESS);
    EXPECT_EQ(aclopExecWithHandle(handle, 2, inputs_, 1, outputs_, stream), ACL_SUCCESS);

    free(data);
}


TEST_F(OpExecutorTest, TestDynamicExecuteAsyncWithOpHandle)
{
    EXPECT_CALL(MockFunctionTest::aclStubInstance(), LoadDynamicSingleOpV2(_,_,_,_,_))
        .WillOnce(Return(ACL_ERROR_GE_LOAD_MODEL))
        .WillRepeatedly(Return(SUCCESS));
    
    EXPECT_CALL(MockFunctionTest::aclStubInstance(), ExecuteAsync(_,_,_,_,_))
        .WillRepeatedly(Return(SUCCESS));

    aclopHandle *handle = nullptr;

    OpModelDef modelDef;
    modelDef.opType = "dynamic_op";
    modelDef.modelPath = "BBB";

    int64_t shape[] = {-1, -1};
    aclTensorDesc desc(ACL_FLOAT, 2, shape, ACL_FORMAT_ND);
    desc.shapeRange.emplace_back(std::pair<int64_t, int64_t>(1, -1));
    desc.shapeRange.emplace_back(std::pair<int64_t, int64_t>(1, -1));

    modelDef.inputDescArr.emplace_back(desc);
    modelDef.inputDescArr.emplace_back(desc);
    modelDef.outputDescArr.emplace_back(desc);

    OpModel opModel;

    char_t *data = new char_t(32);
    const std::shared_ptr<void> p_data(data, [](const char_t *const p) { delete[]p; });
    opModel.data = p_data;
    opModel.size = 32U;
    opModel.name = "path";

    auto &instance = OpModelManager::GetInstance();
    instance.modelCache_.Add(modelDef, opModel);
    EXPECT_EQ(instance.RegisterModel(std::move(modelDef), OpModelManager::GetInstance().opModels_, true), ACL_SUCCESS);

    aclopAttr *attr = aclopCreateAttr();
    aclopCreateHandle("dynamic_op", 2, input_desc_, 1, output_desc_, attr, &handle);
    const aclrtStream stream = (aclrtStream) 0x1234;
    EXPECT_NE(aclopExecWithHandle(handle, 2, inputs_, 1, outputs_, stream), ACL_SUCCESS);
    EXPECT_EQ(aclopExecWithHandle(handle, 2, inputs_, 1, outputs_, stream), ACL_SUCCESS);
    
    instance.modelCache_.Delete(modelDef, false);
    free(data);
}
